import hashlib, os, sqlite3,time,sys
from os import cpu_count
from random import random

from mainui import *
from PyQt6.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox,QTreeWidgetItem
from PyQt6.QtCore import pyqtSlot,Qt,pyqtSignal, QThread
from concurrent.futures import ThreadPoolExecutor
import multiprocessing
import tempfile

def multi_md5(path_to_db, path_to_files, queue):

    file_count = 0
    # 连接数据库
    conn = sqlite3.connect(path_to_db)
    with open(path_to_files,'r', encoding='utf-8') as file:
        # 读取第一行
        path_to_md5 = file.readline().strip()
        while path_to_md5 != '':
            file_count += 1
            # queue 发送信号
            queue.put(path_to_md5)

            # =============
            # print(path_to_md5)
            # ================

            # print('\n', path_to_md5, end='', flush=True)
            # 获取文件的md5 储存到sqlite3数据库中

            with open(path_to_md5, 'rb') as f:
                md5 =  hashlib.md5(f.read()).hexdigest()
            volume = os.path.getsize(path_to_md5)
            file_name = os.path.split(path_to_md5)[-1]
            suffix = os.path.splitext(file_name)[-1][1:]
            conn.execute(
                '''INSERT INTO FILEMAP (filepath,md5,filename,suffix,volume) VALUES ("{}", "{}", "{}","{}", {})'''.format(
                    path_to_md5, md5,
                    file_name, suffix, volume))
            # 提交到数据库
            conn.commit()
            # 读取下一行
            path_to_md5 = file.readline().strip()
    # 关闭数据库
    conn.close()
    return file_count


class ui_update(QThread):
    info_signal = pyqtSignal(str)
    compare_signal = pyqtSignal()
    def __init__(self, handle, queue):
        super().__init__()
        self.handle = handle
        self.queue = queue
        self.files_num = None
        self.files_count = None
        self.quit_signal = False
        self.timer_start = 0
        self.timer_end = 0
        # print(self.handle)
    def run(self):
        self.timer_start = time.time()
        self.files_count = 0
        while True:
            if not self.queue.empty():
                self.files_count += 1
                text = '{}% - {}/{} {}'.format(round(((self.files_count * 100) / self.files_num), 2),
                                                 self.files_count, self.files_num, self.queue.get())
                self.info_signal.emit(text)
            if self.files_count == self.files_num:
                print("对比文件差异")
                self.info_signal.emit("正在对比文件差异......")
                self.compare_signal.emit()
                self.timer_end = time.time()
                print("用时{}".format(round(self.timer_end - self.timer_start,2)))
                break



class MainWindow(QMainWindow):
    finished_signal = pyqtSignal()
    def __init__(self):
        super().__init__()
        self.processing_file = None
        self.split_sheet = None
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        # 文件计数器
        self.file_count = 0
        self.file_num = 0
        # 核心数
        self.cpus = multiprocessing.cpu_count() // 2
        # self.cpus = 1
        # 多进程queue
        self.queue = multiprocessing.Manager().Queue()
        # 工作线程
        # 使用Qthread # 任务线程
        self.worker = ui_update(self, self.queue)
        self.worker.info_signal.connect(self.update_ui)     # 连接到ui更新函数
        self.worker.compare_signal.connect(self.compare_md5)# 连接到对比完成的函数

    @pyqtSlot()
    def on_start_clicked(self):
        # 改为直接在主进程调用multiprocessing.Pool
        self.traverse(self.ui.path.text())
        # 把文件总数传递给工作进程
        self.worker.files_num = self.file_num
        self.worker.start()


    @pyqtSlot()
    def on_onlyone_clicked(self):
        count = self.ui.tree.topLevelItemCount()
        for i in range(count):
            mainitem = self.ui.tree.topLevelItem(i)
            childcount = mainitem.childCount()
            for n in range(childcount):
                item = mainitem.child(n).setCheckState(0,Qt.CheckState.Unchecked)
            mainitem.child(0).setCheckState(0,Qt.CheckState.Checked)
    @pyqtSlot()
    def on_choose_clicked(self):
        path = QFileDialog.getExistingDirectory(self,"选择一个目录")
        if os.path.isdir(path):
            self.ui.path.setText(path)
    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
        try:
            if os.path.isfile(self.dbfile):
                os.remove(self.dbfile)
        except Exception as e:
            pass
    # 一对多字典{"G://1//1.jpg":{"md5":"156fsdf6as4e5","name":"1.jpg","suffix":"jpg","volumn":1565}
    def update_ui(self, text):
        # 更新ui
        # print(text)
        self.ui.label_2.setText(text)
    def get_md5(self,path):
        with open(path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()
    def hidebutton(self):
        # 隐藏按钮
        if self.ui.tree.topLevelItemCount() == 0:
            self.ui.remove.setEnabled(False)
            self.ui.keep.setEnabled(False)
            self.ui.onlyone.setEnabled(False)
    def database(self):
        # 临时文件夹
        temp_folder = tempfile.gettempdir()
        print(temp_folder)
        # 创建或连接数据库
        # db_path = "{}/filemap_{}.db".format(temp_folder,time.strftime("%Y%m%d%H%M%S"))
        db_path = "{}/{}.db".format(temp_folder, 'DupulyMD5')
        if os.path.exists(db_path):
            os.remove(db_path)
        self.dbfile = db_path
        conn = sqlite3.connect(self.dbfile, )
        #启用wal模式减缓并发崩溃现象
        conn.execute('PRAGMA journal_mode=WAL')
        conn.execute('''
        CREATE TABLE FILEMAP(
        	filepath text,
        	md5 text,
        	filename text,
        	suffix text,
        	volume int
        )
        ''')
        conn.commit()
        # 关闭
        conn.close()
        return conn



    def traverse(self,index):
        self.ui.start.setEnabled(False)
        self.ui.path.setEnabled(False)
        self.ui.choose.setEnabled(False)
        # 比对过程中先隐藏按钮
        self.ui.remove.setEnabled(False)
        self.ui.keep.setEnabled(False)
        self.ui.onlyone.setEnabled(False)
        # 初始化
        self.file_num = 0
        self.file_count = 0
        self.processing_file = []
        # 创建数据库文件
        conn = self.database()

        # 计算总共文件数量
        DupulyFilesList = tempfile.gettempdir()+'DupulyFilesList.txt'
        print(DupulyFilesList)
        with open(DupulyFilesList,'w', encoding="utf-8") as file:
            for prefix, _, files in os.walk(index):
                for filename in files:
                    file.writelines(str(os.path.join(prefix, filename)) + '\n')
                    self.file_num += 1
        base_split_num = self.file_num // self.cpus
        remainder = self.file_num % self.cpus
        self.split_sheet = [base_split_num + (i < remainder) for i in range(self.cpus)]
        # print(self.split_sheet)

        # 分割成多个list文件，用于多线程任务
        with open(DupulyFilesList, "r", encoding="utf-8") as file:
            # splite_sheet 非零数量
            splite_sheet_zero = self.split_sheet.count(0)
            splite_sheet_none_zero = len(self.split_sheet) - splite_sheet_zero
            for i in range(splite_sheet_none_zero):
                processing_file = tempfile.gettempdir() + 'DupulyFilesList_{:>03d}.txt'.format(i)
                self.processing_file.append(processing_file)
                with open(processing_file, 'w', encoding="utf-8") as th:
                    for j in range(self.split_sheet[i]):
                        th.write(file.readline())
        # print(self.processing_file)

        print("总共{}个文件".format(self.file_num))
        # 遍历文件夹，计算md5 通过线程池进行多线程的md5扫描和储存工作
        # 使用with上下文管理器，相当于自动调用了executor.shutdown(wait=True)
        # 会等待线程池的线程结束再进入主进程后续代码
        # with ThreadPoolExecutor(max_workers=10) as executor:
        #     for prefix, _, files in os.walk(index):\
        #         # 把当前目录下每一个文件记录
        #         for file in files:
        #             executor.submit(self.storge_md5,conn, file, prefix)
        # 显式调用pool
        pool = multiprocessing.Pool(self.cpus)
        self.tasks = []
        for file in self.processing_file:
            # print("process 启动 {}".format(file))
            self.tasks.append(pool.apply_async(multi_md5, (self.dbfile, file, self.queue)))
        pool.close()
        # pool.join()

    def compare_md5(self):
        result = sum([task.get() for task in self.tasks])
        # print([task.get() for task in tasks])
        # print("\n{}/{}".format(self.file_count, self.file_num))
        # print("\n{}/{}".format(result, self.file_num))
        # 对比md5
        conn = sqlite3.connect(self.dbfile)
        self.searchdu(conn)
        self.hidebutton()
        conn.close()
        try:
            if os.path.isfile(self.dbfile):
                os.remove(self.dbfile)
        except Exception as e:
            pass
        self.ui.start.setEnabled(True)
        self.ui.path.setEnabled(True)
        self.ui.choose.setEnabled(True)
    def test(self,a,b,c):
        print(a,b,c)
    def storge_md5(self, conn, file, prefix):
        print("===============")
        # 连接数据库
        # conn = sqlite3.connect(self.dbfile)
        # 增加一个文件计数
        self.file_count += 1
        filepath = os.path.join(prefix, file)
        print('\r', filepath, end='', flush=True)
        self.ui.label_2.setText("{}%-".format(round((self.file_count*100) / self.file_num,2)) + str(filepath))
        # 获取文件的md5 储存到sqlite3数据库中
        md5 = self.get_md5(filepath)
        volume = os.path.getsize(filepath)
        suffix = os.path.splitext(file)[-1][1:]
        conn.execute(
            '''INSERT INTO FILEMAP (filepath,md5,filename,suffix,volume) VALUES ("{}", "{}", "{}","{}", {})'''.format(
                filepath, md5,
                file, suffix, volume))
        # 提交到数据库
        conn.commit()

        # self.searchdu(conn)
        # self.hidebutton()
        # conn.close()
        # try:
        #     if os.path.isfile(self.dbfile):
        #         os.remove(self.dbfile)
        # except Exception as e:
        #     pass
        # self.ui.start.setEnabled(True)
        # self.ui.path.setEnabled(True)
        # self.ui.choose.setEnabled(True)

    def searchdu(self,conn):
        # print("开始比对")
        self.ui.label_2.setText("开始比对")
        self.ui.tree.clear()
        cur = conn.cursor()
        # 根据md5 来查重
        query = '''SELECT * FROM FILEMAP WHERE {} in\
         (SELECT {} FROM FILEMAP GROUP BY {} HAVING COUNT({}) >=2) ORDER BY {}'''.replace("{}", "md5")
        # query='''SELECT * FROM FILEMAP WHERE filepath in\
        #  (SELECT filepath FROM FILEMAP GROUP BY volume HAVING COUNT(volume) >=2)'''
        cur.execute(query)
        data = cur.fetchall()
        files_num = len(data)
        cur.execute('''SELECT * FROM FILEMAP GROUP BY {} HAVING COUNT({}) >=2 ORDER BY {}'''.replace("{}", "md5"))
        md5 = cur.fetchall()
        unique_file = len(md5)
        # print(len(data))
        # 添加父节点
        for md5_p in md5:
            mainitem = QTreeWidgetItem(0)
            mainitem.setText(0, "文件大小：{}  文件类型：{}  文件md5：{}".format(str(md5_p[4]),md5_p[3],md5_p[1]))
            self.ui.tree.addTopLevelItem(mainitem)
            # self.ui.tree.setCurrentIndex(-1)
            for i in data:
                if i[1] == md5_p[1]:
                    # print(i)
                    # 添加行
                    item = QTreeWidgetItem(1)
                    item.setText(0,i[0])
                    item.setText(1,i[2])
                    item.setText(2,i[1])
                    item.setText(3,i[3])
                    item.setText(4,str(i[4]))
                    # item.setFlags((Qt.ItemIsSelectable | Qt.ItemIsUserCheckable  | Qt.ItemIsEnabled | Qt.ItemIsAutoTristate))
                    item.setCheckState(0,Qt.CheckState.Unchecked)
                    mainitem.addChild(item)
        info = '''共检查到{}项可能重复文件，根据md5校验出{}项独特的文件'''.format(files_num,unique_file)
        self.ui.label_2.setText(info)
        if self.ui.tree.topLevelItemCount():
            self.ui.tree.expandAll()
            self.ui.remove.setEnabled(1)
            self.ui.keep.setEnabled(1)
            self.ui.onlyone.setEnabled(1)


    @pyqtSlot()
    def on_remove_clicked(self):
        if QMessageBox.question(self, "重要提示！", "是否确定删除所有选择的文件，请仔细核对。该操作无法撤回！！",
                                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.Yes) == QMessageBox.StandardButton.No:
            QMessageBox.information(self,"通知", "取消操作")
            return 0
        remove_list = []
        count = self.ui.tree.topLevelItemCount()
        offset = 0
        for i in range(count):
            mainitem = self.ui.tree.topLevelItem(i-offset)
            self.deeptree_remove(mainitem,remove_list)
            if mainitem.childCount() <= 1:
                # 子节点小于等于1，清除父节点
                self.ui.tree.takeTopLevelItem(i-offset)
                offset+=1

        print(remove_list)
        self.ui.label_2.setText('''已删除{}项重复文件'''.format(len(remove_list)))
        self.hidebutton()

    def deeptree_remove(self,mainitem,remove_list):
        count = mainitem.childCount()
        offset = 0
        for i in range(count):
            item = mainitem.child(i-offset)
            if item.checkState(0) == Qt.CheckState.Checked:
                offset+=1
                os.remove(item.text(0))
                item.parent().removeChild(item)
                remove_list.append(item.text(0))

    def deeptree_keep(self,mainitem,remove_list):
        count = mainitem.childCount()
        offset = 0
        for i in range(count):
            item = mainitem.child(i-offset)
            if item.checkState(0) == Qt.CheckState.Unchecked:
                offset+=1
                os.remove(item.text(0))
                item.parent().removeChild(item)
                remove_list.append(item.text(0))

    @pyqtSlot()
    def on_keep_clicked(self):
        if QMessageBox.question(self, "重要提示！", "是否确定仅保留所选文件，请仔细核对。该操作无法撤回！！",
                                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.Yes) == QMessageBox.StandardButton.No:
            QMessageBox.information(self,"通知", "取消操作")
            return 0
        remove_list = []
        count = self.ui.tree.topLevelItemCount()
        offset = 0
        for i in range(count):
            mainitem = self.ui.tree.topLevelItem(i-offset)
            self.deeptree_keep(mainitem, remove_list)
            if mainitem.childCount() <= 1:
                # 子节点小于等于1，清除父节点
                self.ui.tree.takeTopLevelItem(i-offset)
                offset += 1
        print(remove_list)
        self.ui.label_2.setText('''已删除{}项重复文件'''.format(len(remove_list)))
        self.hidebutton()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    mywin = MainWindow()
    mywin.show()
    sys.exit(app.exec())